import cmd
import sys
import os
import functools
import ollama
import json
import fitz  # PyMuPDF


# 使用base code中的common 和 pocketflow
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)
from base_code import *
from base_code.common import *

# ==============================================================================
# simple_mas_shell 代码介绍
# 在这个案例中，我们引入两个Agent分别处理命令执行和界面显示（交互逻辑仍由cmd库处理）
#  - BDIAgent： 构建命令的工作流
#  - HelpAgent： 显示帮助窗口
#  - CommandRegistry： 存储当前支持的命令信息
# 我们会继续拓展这个案例，来看如何简化这个交互，使用户能充分利用所有已实现工具
# ==============================================================================

# ==============================================================================
# 1. CommandRegistry 是ToolSystem的继承类，添加对支持的命令列表的显示
# ==============================================================================
class CommandRegistry(ToolSystem):
    def __init__(self):
        super().__init__()
        self.lines = [
            "     欢迎使用 SimpleMAS plus v0.3 (工作流版)",
            " ",
            "支持自然语言交互，所有逻辑由Agent的P-D-A循环由工作流处理",
            " ",
            "set <key> <value>        - 设置一个键值对",
            "get <key>                - 获取一个键对应的值",
            "delete <key>             - 删除一个键值对",
            "list                     - 列出所有存储的键值对",
            "use_tool <tool> <param>  - 使用工具",
            "ls path                  - 列出path目录中的文件列表",
            "readpdf <pdf>            - 列出path目录中的文件列表",
            "help                     - 显示帮助信息",
            "exit                     - 退出程序"
        ]


##################################################################
# 全局单例对象：工具系统（环境层）
##################################################################
GLOBAL_TOOL_SYSTEM = CommandRegistry()


##################################################################
# 装饰器：简化工具注册
##################################################################
def tool(name: str, description: str, required_args: List[str], optional_args: Optional[Dict[str, Any]] = None):
    """
    用于注册外部工具函数的装饰器。
    - name: 工具名
    - description: 工具描述
    - required_args: 必须参数列表
    - optional_args: 可选参数字典
    """
    def decorator(func: Callable):
        # 将原函数包装为 Tool 并注册
        new_tool = Tool(
            name=name,
            description=description,
            function=func,
            required_args=required_args,
            optional_args=optional_args
        )
        GLOBAL_TOOL_SYSTEM.register_tool(new_tool)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapper
    return decorator


# ==============================================================================
# 2. HelpAgent 支持UI绘制
# ==============================================================================
class HelpAgent:
    def __init__(self, registry: CommandRegistry):
        self.command_registry = registry

    def show_intro(self):
        try:
            return create_aligned_intro(self.command_registry.lines)
        except Exception:
            return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(t.required_args)}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()


# ==============================================================================
# 3. BDIAgent 的核心组件
# ==============================================================================
class AgentMemory:
    """
    Memory class that stores short and long memory of agents

    """
    def __init__(self):
        self.beliefs = {}
        self.raw_input = None
        self.command = None
        self.raw_arg = None
        self.desire = None       # (tool_name, kwargs_dict)
        self.error_message = None
        self.result = None
        self.registry = None

    def reset_workflow_data(self):
        """在每次运行前重置临时数据。"""
        self.command, self.raw_arg, self.desire, self.error_message, self.result = (None,) * 5


class PerceptionNode(AsyncNode):
    """
    感知节点（自然语言版），负责解析用户输入并生成执行愿望(desire)。

    本节点采用混合策略：
    1.  **快速路径**：首先尝试将输入作为直接命令进行解析。如果输入以一个已注册的工具名开头，
        则直接构建desire，跳过LLM调用，以提高效率和确定性。
    2.  **慢速路径**：如果输入不是一个直接命令，则调用已注册的 'ollama_call' 工具，
        利用大型语言模型将自然语言转换为结构化的命令。

    输入: shared.raw_input (用户的原始输入字符串)
    输出:
      - 成功时: shared.desire 被设置为 (tool_name, kwargs)，并返回 "success"。
      - 失败时: shared.error_message 被设置，并返回 "error"。
    """

    def __init__(self, model_name: str = "gemma3:4b"):
        """
        初始化感知节点。
        Args:
            model_name (str): 用于自然语言解析的Ollama模型名称。
        """
        super().__init__()
        self.model_name = model_name
        # 将系统提示词的生成逻辑放在初始化中，避免重复构建
        self.system_prompt = self._get_system_prompt()

    async def prep_async(self, shared: AgentMemory) -> AgentMemory:
        return shared
    def _get_system_prompt(self) -> str:
        """
        生成用于指导LLM进行命令解析的系统提示词。
        这是实现自然语言理解的关键部分。
        """
        tools_description = """
            可用的命令如下:
            - 'set <key> <value>': 设置一个键值对。值可以包含空格。
            - 'get <key>': 获取一个键对应的值。
            - 'delete <key>': 删除一个键值对。
            - 'list': 列出所有存储的键值对，此命令没有参数。
            - 'ls <path>': 列出目录中的文件列表。如果未提供路径，则使用当前目录。
            - 'list_tools': 列出ToolSystem中注册的工具。
            - 'read_pdf' <path> <max_page> : 从path给出的pdf读取最多max_page页内容
            - 'quit': 退出系统。
            - 'unknown': 如果用户意图不明确或与任何可用命令都不匹配。
            """

        return f"""
            你是一个智能助手，你的唯一任务是将用户的自然语言请求转换为一个精确的JSON对象。
            这个JSON对象必须包含两个键：'command' 和 'arguments'。
            'command' 的值必须是以下之一：'set', 'get', 'delete', 'list', 'ls', 'list_tools', 'quit', 'unknown'。
            'arguments' 的值必须是一个字符串列表。

            {tools_description}

            根据用户的请求，严格按照以下JSON格式进行响应，不要添加任何额外的解释或文本。

            例如:
            - 用户输入: "帮我把 'user:id' 设置成 '12345'"
              你的输出: {{"command": "set", "arguments": ["user:id", "12345"]}}
            - 用户输入: "查一下 'user:id' 的值"
              你的输出: {{"command": "get", "arguments": ["user:id"]}}
            - 用户输入: "显示所有东西"
              你的输出: {{"command": "list", "arguments": []}}
            - 用户输入: "列出 C:\\Users 目录下的文件"
              你的输出: {{"command": "ls", "arguments": ["C:\\\\Users"]}}
            - 用户输入: "今天天气怎么样"
              你的输出: {{"command": "unknown", "arguments": ["无法处理天气查询"]}}
            - 用户输入：读取当前目录下的'paper.pdf'并分析"
              你的输出: {{"command": "read_pdf", "arguments": [paper.pdf"]}}
      
             """

    async def exec_async(self, prep_res):
        user_input = (prep_res.raw_input or "").strip()
        registry = prep_res.registry

        if not user_input:
            # 如果没有输入，则什么都不做，避免错误
            prep_res.error_message = "ℹ️ 请输入指令。"
            return "error"

        if not registry:
            prep_res.error_message = "❌ Agent 未连接到 registry"
            return "error"

        # --- 1. 快速路径：尝试直接匹配命令 ---
        parts = user_input.split(None, 1)
        potential_cmd = parts[0] if parts else ""

        if registry.has_tool(potential_cmd):
            print("⚡️ 快速路径：识别到直接命令。")
            prep_res.command = potential_cmd
            prep_res.raw_arg = parts[1] if len(parts) > 1 else ""

        # --- 2. 慢速路径：如果不是直接命令，则调用LLM ---
        else:
            print("🧠 慢速路径：调用LLM进行自然语言理解...")
            try:
                # 动态获取并调用'ollama_call'工具
                ollama_tool = registry.get_tool("ollama_call")
                nlp_result = await ollama_tool(
                    memory=prep_res,
                    model=self.model_name,
                    system_prompt=self.system_prompt,
                    user_prompt=user_input,
                    as_json=True  # 确保返回的是字典
                )
            except KeyError:
                prep_res.error_message = "❌ 内部错误: 'ollama_call' 工具未在系统中注册。"
                return "error"
            except Exception as e:
                prep_res.error_message = f"❌ 调用LLM工具时发生意外错误: {e}"
                return "error"

            # --- 3. 处理LLM的返回结果 ---
            command = nlp_result.get('command')
            arguments = nlp_result.get('arguments', [])
            print(f"command: {command}, arguments: {arguments}")

            if command == 'error' or not command:
                error_details = ' '.join(map(str, arguments))
                prep_res.error_message = f"❌ NLP 处理失败: {error_details or '模型返回无效响应。'}"
                return "error"

            if command == 'unknown':
                prep_res.error_message = f"🤔 无法理解您的指令 '{user_input}'。请尝试与 'set', 'get', 'list', 'delete', 'ls' 相关的操作。"
                return "error"

            prep_res.command = command
            # 将LLM返回的参数列表重新组合成字符串，以适应后续工具的'raw_arg'输入
            prep_res.raw_arg = " ".join(map(str, arguments))

        # --- 4. 统一的Desire生成 ---
        # 无论通过快速还是慢速路径，此时都应该有 shared.command
        if not prep_res.command or not registry.has_tool(prep_res.command):
            prep_res.error_message = f"❌ 解析后的命令 '{prep_res.command}' 不是一个有效的工具。"
            return "error"

        # 构建desire，供ToolExecutionNode使用
        kwargs = {"memory": prep_res, "raw_arg": prep_res.raw_arg}
        prep_res.desire = (prep_res.command, kwargs)


        print(f"✅ 准备执行命令 '{prep_res.command}', 参数: '{prep_res.raw_arg}'")
        return "success"

    async def post_async(self, shared, prep_res, exec_res):
        return exec_res


class DeliberationNode(Node):
    """
    决策节点：根据愿望选择一个计划（意图），返回对应的动作。
    """
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        # 决策简单：desire 已经包含目标工具与参数，直接走成功路径
        return "success"

    def post(self,shared,prep_res, exec_res):
        return exec_res


class ToolExecutionNode(AsyncNode):
    def __init__(self, tool_system: ToolSystem):
        super().__init__()
        self.tool_system = tool_system

    async def prep_async(self, shared: AgentMemory) -> AgentMemory:
        return shared

    async def exec_async(self, shared: AgentMemory):
        if not shared.desire:
            shared.result = "❌ 无欲望(desire)来执行"
            return

        tool_name, kwargs = shared.desire
        try:
            tool_func = self.tool_system.get_tool(tool_name)
        except KeyError:
            shared.result = f"❌ 工具 '{tool_name}' 未找到"
            return

        # 确保 memory 注入
        if "memory" not in kwargs:
            kwargs["memory"] = shared

        try:
            maybe_coro = tool_func(**kwargs)
            if asyncio.iscoroutine(maybe_coro):
                ret = await maybe_coro
            else:
                ret = maybe_coro
            # 工具函数可以直接把结果写入 memory.result，也可以返回字符串
            if isinstance(ret, str) and (not shared.result):
                shared.result = ret
        except Exception as e:
            shared.result = f"❌ 工具 '{tool_name}' 执行失败: {e}"

    def post(self,shared,prep_res, exec_res):
        return exec_res


class ErrorNode(Node):
    """处理流程中产生的错误。"""
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        shared.result = shared.error_message


class FinalizeNode(Node):
    """
    所有成功路径的终点，确保结果被正确设置。
    这个节点目前什么也不做，只是作为一个清晰的汇聚点。
    """
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        pass


# ==============================================================================
#  基于流程的BDIAgent定义
# ==============================================================================
class BDIAgent(AsyncFlow):
    """
    一个本身就是异步工作流的BDI Agent。
    它继承自AsyncFlow，其流程定义即其内部逻辑。
    """
    def __init__(self):
        # 首先调用父类的构造函数
        super().__init__()

        # 集成 memory 部件
        self.memory = AgentMemory()
        # 引用全局GLOBAL_TOOL_SYSTEM（ToolSystem)
        self.memory.registry = GLOBAL_TOOL_SYSTEM

        # 1. 定义nodes
        perception = PerceptionNode()
        deliberation = DeliberationNode()
        error_handler = ErrorNode()
        finalizer = FinalizeNode()
        executor = ToolExecutionNode(self.memory.registry)

        # 2. 编排工作流 (使用 self.start() 和 self.next())
        self.start(perception)

        perception - "success" >> deliberation
        perception - "error" >> error_handler
        deliberation - "success" >> executor
        executor >> finalizer

        print("💡 BDIAgent (Inheritance-based) 已初始化。Agent本身即是工作流。")

    async def process_command_async(self, line: str) -> str:
        """
        Agent的异步入口方法。

        :param command: 命令名
        :param raw_arg: 原始参数字符串
        :return: 执行结果
        """
        # 1. 准备记忆，理论上应放到belief中(sense)，但此处belief主要用于存储用户输入键值
        self.memory.reset_workflow_data()
        self.memory.raw_input = line

        # 2. 执行自身所代表的异步工作流
        #    perception -> ...
        await self.run_async(self.memory)

        # 3. 返回结果
        return self.memory.result



# ==============================================================================
# 注册工具
# ==============================================================================
@tool("set",  "设置键值", ["memory", "raw_arg"])
async def set_value(memory: AgentMemory, raw_arg: str):
    parts = raw_arg.split(None, 1)
    if len(parts) < 2:
        return "❌ 错误: 'set' 命令需要一个键和一个值。\n   用法: set <key> <value>"
    key, value = parts
    memory.beliefs[key] = value
    return f"✅ '{key}' 已设置为 '{value}'"

@tool("get",  "获取键值", ["memory", "raw_arg"])
async def get_value(memory: AgentMemory, raw_arg: str):
    key = raw_arg.strip()
    if not key:
        return "❌ 错误: 'get' 需要一个键。 用法: get <key>"
    if key in memory.beliefs:
        return f"➡️ {key} = {memory.beliefs[key]}"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("list",  "列出键值", ["memory", "raw_arg"])
async def list_beliefs(memory: AgentMemory, raw_arg: str):
    if not memory.beliefs:
        return "ℹ️ 数据库为空"
    else:
        return "\n".join([f"{k}: {v}" for k, v in memory.beliefs.items()])


@tool("delete",  "删除键值", ["memory", "raw_arg"])
async def delete_key(memory: AgentMemory, raw_arg: str):
    key = raw_arg.strip()
    if not key:
        return "❌ 错误: 'delete' 需要一个键。 用法: delete <key>"
    if key in memory.beliefs:
        del memory.beliefs[key]
        return f"✅ 成功: 键 '{key}' 已被删除。"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("ls",  "列出目录", ["memory","raw_arg"])
async def ls_dir(memory: AgentMemory, raw_arg: str = '.'):
    path = raw_arg.strip() or "."
    try:
        if not os.path.isdir(path): return f"错误: '{path}' 不是一个有效的目录。"
        files = os.listdir(path)
        return f"目录 '{path}' 下的内容:\n- " + "\n- ".join(files) if files else f"目录 '{path}' 为空。"
    except Exception as e:
        return f"执行时发生错误: {e}"


@tool("use_tool",  "使用工具", ["memory","raw_arg"])
async def use_tool(memory: AgentMemory, raw_arg: str ):
    arg = raw_arg.strip()
    if not arg:
        return "❌ 用法: use_tool <tool_name> [args]"

    parts = arg.split(None, 1)
    tool_name = parts[0]
    inner_raw = parts[1] if len(parts) > 1 else ""

    if not memory.registry.has_tool(tool_name):
        memory.result = f"❌ 没有找到工具 '{tool_name}'"
        return "error"

    tool_func = GLOBAL_TOOL_SYSTEM.get_tool(tool_name)
    # 调用目标工具，注入 memory 和 inner_raw
    try:
        maybe = tool_func(memory=memory, raw_arg=inner_raw)
        if asyncio.iscoroutine(maybe):
            return await maybe
    except Exception as e:
        memory.result = f"❌ 调用工具失败: {e}"
        return "error"


@tool("ollama_call", "调用本地Ollama LLM", ["memory", "model", "system_prompt", "user_prompt"], ["as_json"])
async def ollama_call(memory: AgentMemory, model: str, system_prompt: str, user_prompt: str, as_json: bool = True):
    try:
        response = await ollama.AsyncClient().chat(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            format="json" if as_json else None
        )

        content = response["message"]["content"]

        if as_json:
            try:
                return json.loads(content)
            except json.JSONDecodeError:
                return {"command": "error", "arguments": ["无法解析为JSON", content]}
        else:
            return content

    except ollama.ResponseError as e:
        return {"command": "error", "arguments": [f"Ollama API 错误: {e.error}"]}
    except Exception as e:
        return {"command": "error", "arguments": [f"未知错误: {str(e)}"]}


@tool("read_pdf", "读取pdf", ["memory", "raw_arg"])
def read_pdf_text(memory: AgentMemory, raw_arg: str, max_pages: int = 1) -> str:
    """
    读取 PDF 前若干页的文本内容
    """
    key = raw_arg.split(None, 1)
    if not key:
        return "❌ 错误: 'read_pdf' 需要pdf文件路径"
    path = key[0]
    if not os.path.exists(path):
        raise FileNotFoundError(f"PDF file not found: {path}")

    if not path.lower().endswith('.pdf'):
        raise ValueError("文件必须是PDF格式")

    text = []
    try:
        with fitz.open(path) as doc:
            # 基础信息
            info = {
                "path": path,
                "pages": doc.page_count,
                "metadata": doc.metadata,
                "chapters": []
            }

            for page_num in range(min(len(doc), max_pages)):
                page = doc.load_page(page_num)
                text.append(page.get_text())
        return "\n".join(text)
    except Exception as e:
        return f"执行时发生错误: {e}"



@tool("search_pdf", "搜索pdf", ["memory", "path", "keyword"])
def search_pdf(memory: AgentMemory, path: str, keyword: str) -> list:
    """
    在 PDF 中搜索关键字，返回匹配的页码和上下文。
    """
    if not os.path.exists(path):
        raise FileNotFoundError(f"PDF file not found: {path}")

    results = []
    with fitz.open(path) as doc:
        for page_num, page in enumerate(doc, start=1):
            text = page.get_text()
            if keyword.lower() in text.lower():
                snippet = text[text.lower().index(keyword.lower())-30:
                              text.lower().index(keyword.lower())+30]
                results.append({"page": page_num, "context": snippet})

    return results
# =======================================================================
# SimpleAgentShell（应用层）
# =======================================================================
class ReActTerminal(cmd.Cmd):
    """
    交互式Shell，通过 asyncio.run() 调用Agent的异步方法完成用户指令。
    """
    help_agent = HelpAgent(GLOBAL_TOOL_SYSTEM)
    cmd_agent = BDIAgent()
    intro = help_agent.show_intro()
    prompt = "agent> "

    def __init__(self):
        super().__init__()

    def default(self, line):
        """
        所有未知命令，直接送到 Agent
        """
        if not line.strip():
            return
        result = asyncio.run(self.cmd_agent.process_command_async(line))
        print(result or (self.cmd_agent.memory.result if self.cmd_agent.memory else ""))

    def do_list_tools(self, arg):
        print("可用命令:", self.cmd_agent.memory.registry.list_tools())

    # --- UI控制相关方法
    def do_exit(self, arg):
        print("\n👋 感谢使用，再见！")
        return True
    def do_quit(self, arg): return self.do_exit(arg)
    def do_EOF(self, arg): print(); return self.do_exit(arg)
    def do_help(self, arg):
        print(self.help_agent.show_help(arg.strip()))

def main():
    """程序主入口。"""
    try:
        ReActTerminal().cmdloop()
    except KeyboardInterrupt:
        print("\n👋 检测到中断，程序退出。")
        sys.exit(0)

if __name__ == '__main__':
    main()